#!/usr/bin/env python

from numpy import array, pi, sqrt, zeros


'''
Each method in this module reads in a text file in the profiles directory:
'''

def load_clusterheader( line ) :
    '''
    This parses through the hblist line before the profiles
    # Header Columns:
    # id min(rout,rmax) min(rvir,rmax) [/h kpc]
    # Mg Mcg M* M*new Mb Mdm Mtotal [/h Msolar,< rvir] Mvir (from halo finder)
    # vmax [km/s] rmax [/h kpc]
    # <Zg_II> <Zg_Ia> <Z*_II> <Z*_Ia> <Z*new_II> <Z*new_Ia> <t*> [Gyr]
    
    hblist = { 'id':{}, 'rout':{}, 'rvir':{},
    'Mg':{}, 'Mcg':{}, 'Mst':{}, 'Mnst':{},
    'Mb':{}, 'Mdm':{}, 'Mtot':{}, 'Mvir':{},
    'vmax':{}, 'rmax':{},
    'Zg_II':{}, 'Zg_Ia':{}, 'Zst_II':{}, 'Zst_Ia':{},
    'Zstnew_II':{}, 'Zstnew_Ia':{}, 'tst':{}
    }
    Ex:
    hblist=load_clusterheader( line_from_r500_header )
    Mg500=hblist['Mg'] # Returns the gas mass enclosed in r500
    '''
    hblist = { }             

    if not line.startswith('# ') or len(line.split()) != 21 :
        print 'Error: Not a clusterheader line, check binning:\n',line
        print 'Header length should be 21. Length is',len(line.split())
        print 'Check analysis.c to make sure printed header is correct'
    
    hblist['id'] = int(line.split()[1])
    
    ( hblist['rout'], hblist['rvir'],
      hblist['Mg'], hblist['Mcg'], hblist['Mst'], hblist['Mnst'],
      hblist['Mb'], hblist['Mdm'], hblist['Mtot'], hblist['Mvir'],
      hblist['vmax'], hblist['rmax'],
      hblist['Zg_II'], hblist['Zg_Ia'], hblist['Zst_II'], hblist['Zst_Ia'],
      hblist['Zstnew_II'], hblist['Zstnew_Ia'], hblist['tst']
      ) = array([ float(x) for x in line.split()[2:] ])    

    return hblist
    

def load_hblist(blist,clusters=None) :
    '''
    Reads in hblist
    Returns dictionary of dictionaries: rout, Mg, Mcg, Mst, Mnst, Mb, Mtotal, vmax, rmax
    Ex: 
    blist=load_hblist(h_blist_a?.????.dat,clusters=[1,2])
    Mg=blist['Mg'][1] # Returns gas mass of halo 1
    '''
    input = open( blist, 'r' )

    # INPUT HEADER #
    # id min(rout,rmax) min(rvir,rmax) [/h kpc]
    # Mg Mcg M* M*new Mb Mdm Mtotal [/h Msolar,< rout] Mvir [/h Msolar]
    # vmax [km/s] rmax [/h kpc]
    # <Zg_II> <Zg_Ia> <Z*_II> <Z*_Ia> <Z*new_II> <Z*new_Ia> <t*> [Gyr]
    
    hblist = {  'rout':{}, 'Mg':{}, 'Mcg':{}, 'Mst':{}, 'Mnst':{}, 'Mb':{},
                'Mdm':{}, 'Mtotal':{}, 'vmax':{}, 'rmax':{}  }

    if ( clusters != None ) :
        cluster_hash = {}
        for c in clusters :
            cluster_hash[c] = 1

    for line in input.readlines() :
        if not line.startswith('#') :
            cols = line.split()
            id = int(cols[0])

            if ( clusters == None or cluster_hash.has_key(id) ) :
                hblist['rout'][id] = float(cols[1])
                hblist['Mg'][id] = float(cols[3])
                hblist['Mcg'][id] = float(cols[4])
                hblist['Mst'][id] = float(cols[5])
                hblist['Mnst'][id] = float(cols[6])
                hblist['Mb'][id] = float(cols[7])
                hblist['Mdm'][id] = float(cols[8])
                hblist['Mtotal'][id] = float(cols[9])
                hblist['vmax'][id] = float(cols[11])
                hblist['rmax'][id] = float(cols[12])
                

    return hblist

def load_bmpro_profile(bmpro,clusters=None,num_bins=100) :
    '''
    Reads in hbmpro (mass profiles), returns density profiles and cumulative mass profiles
    Returns dictionary of dictionaries:
    hbmpro = { 'binr':{}[/h kpc], 'rr':{} [/h kpc], 'Mdm':{}[Msun], 'Mgas':{},
               'Mcg':{}, 'Mst':{}, 'Mnst':{}, 'Mtot':{}, 'header':{} }
    Ex: 
    hbmpro=load_bmpro_profile(h_bmpro_a?.????.dat,clusters=[1,2],num_bins=99)
    Mdm=hbmpro['Mdm'][1] # Returns dm cumulative mass profile (array) of halo 1
    Mdmrout500=hbmpro['header'][1]['Mdm'] # Returns enclosed dm mass within rout of halo 1
    '''

    input = open(bmpro,'r')
    
    if ( clusters != None ) :
        cluster_hash = {}
        for c in clusters :
            cluster_hash[c] = 1
    while 1 :
        line = input.readline()
        if line[0:2] == '##' :
            break

    hbmpro = { 'binr':{}, 'rr':{}, 'Mdm':{}, 'Mg':{},
               'Mcg':{}, 'Mst':{}, 'Mnst':{},
               'Mtot':{}, 'header':{}
               }
    
    while 1 :
        clusterheader = input.readline()
        if ( clusterheader ) :
            header = load_clusterheader( clusterheader )
            id = header['id']
            if ( clusters == None or cluster_hash.has_key(id) ) :
                line = input.readline()
                cols = line.split()

                # Read in first bin of profiles data for this cluster id
                profiles = zeros( (num_bins,len(cols)) )
                profiles[0,:] = array([ float(col) for col in cols ])

                # Read in subsequent bins of profiles data for this cluster id
                for i in xrange(num_bins-1) :
                    line = input.readline()
                    profiles[i+1,:] = array([ float(x) for x in line.split() ])

                ( hbmpro['binr'][id], hbmpro['rr'][id], hbmpro['Mdm'][id],
                  hbmpro['Mg'][id], hbmpro['Mcg'][id],
                  hbmpro['Mst'][id], hbmpro['Mnst'][id]
                  ) = profiles[:,:].transpose()

                hbmpro['Mtot'][id] = hbmpro['Mdm'][id] + hbmpro['Mg'][id]

                # Populate header outputs into hbmpro
                hbmpro['header'][id] = header
                               
                del profiles, header
            else :
                for i in xrange(num_bins) :
                    input.readline()
        else :
            break
    input.close()

    return hbmpro


def load_bgpro_profile(bgpro,clusters=None,num_bins=100) :
    '''
    Reads in h_bgpro (gas property profiles)
    Returns dictionary of dictionaries:
    hbgpro = { 'binr':{}, 'rr':{}, 'Mg':{}, 'Mcg':{},
    'T':{}, 'P':{}, 'K':{},
    'header':{}
    }
    Ex: 
    hbgpro=load_icm_profile(h_bgpro_a?.????.dat,clusters=[1,2],num_bins=99)
    T=hbgpro['T'][1] # Returns temperature profile (array) of halo 1
    '''    
    input = open(bgpro,'r')
    if ( clusters != None ) :
        cluster_hash = {}
        for c in clusters :
            cluster_hash[c] = 1

    while 1 : # Skip over header lines
        line = input.readline()
        
        if line[0:2] == '##' :
            break
    hbgpro = { 'binr':{}, 'rr':{}, 'Mg':{}, 'Mcg':{},
               'T':{}, 'P':{}, 'K':{},
               'header':{}
               }

    while 1 :
        clusterheader = input.readline()
        if ( clusterheader ) :
            header = load_clusterheader( clusterheader )
            id = header['id']

            if ( clusters == None or cluster_hash.has_key(id) ) :
                line = input.readline()
                cols = line.split()
                # Read in first bin of profiles data for this cluster id
                profiles = zeros( (num_bins,len(cols)) )
                profiles[0,:] = array([ float(col) for col in cols ])

                # Read in subsequent bins of profiles data for this cluster id
                for i in xrange(num_bins-1) :
                    line = input.readline()
                    profiles[i+1,:] = array([ float(x) for x in line.split() ])

                ( hbgpro['binr'][id], hbgpro['rr'][id],
                  hbgpro['Mg'][id], hbgpro['Mcg'][id],
                  hbgpro['T'][id], hbgpro['P'][id], hbgpro['K'][id]
                  ) = profiles[:,:].transpose()

                # Populate header outputs into hbgpro
                hbgpro['header'][id] = header
                del profiles, header
            else :
                for i in xrange(num_bins) :
                    input.readline()
        else :
            break
    input.close()

    return hbgpro




def load_bclump_profile(bclumppro,clusters=None,num_bins=100) :
    '''
    Reads in hbclumppro (clumping profiles)
    Returns dictionary of dictionaries:
    keys:
    'binr'[/h kpc], 'rr'[/h kpc], 'Tmw'[K], 'Tsp'[K],
    'rhog'[g/cm^3], 'rhog2_cl'[g^2/cm^6],  dictionary with Tcut as key ('T1','T2','T3','T4')
    'P', 'P_cl' [erg/cm^3],  dictionary with Tcut as key ('T1','T2','T3','T4')
    'mfrac', dictionary with Tcut as key ('T1','T2','T3','T4')
    'header'
    Ex: 
    hclumppro=load_bclump_profile(h_bclumppro_a?.????.dat,clusters=[1,2],num_bins=99)
    rho2cl=hbclumppro['rhog2_cl']['T2'][2] # Returns <rho^2> for gas with T>T2 profile (array) of halo 1
    '''    
    input = open(bclumppro,'r')
    # bclumppro looks like: 
    # T1 T2 T3 T4 = 0.000000e+00 1.000000e+06 5.000000e+06 1.000000e+07 [K]
    # Header Columns:
    # id min(rout,rmax) min(rvir,rmax) [/h kpc]
    # Mg Mcg M* M*new Mb Mdm Mtotal [/h Msolar,< rvir] Mvir (from halo finder)
    # vmax [km/s] rmax [/h kpc]
    # <Zg_II> <Zg_Ia> <Z*_II> <Z*_Ia> <Z*new_II> <Z*new_Ia> <t*> [Gyr]
    # Profile Columns:
    # rmid rr [/h kpc] Tmw Tsp [K] rhog(T1 T2 T3 T4) rhog_cl (T1 T2 T3 T4) [g/cm^3] P(T1 T2 T3 T4) P_cl(T1 T2 T3 T4) [erg/cm^3] mfrac(T1 T2 T3 T4)    

    if ( clusters != None ) :
        cluster_hash = {}
        for c in clusters :
            cluster_hash[c] = 1

    while 1 : # Skip over header lines
        line = input.readline()
        
        if line[0:2] == '##' :
            break
        
    hclumppro = { 'binr':{}, 'rr':{}, 'Tmw':{}, 'Tsp':{},
                  'rhog':{'T0':{},'T1':{},'T2':{},'T3':{},'T4':{}},
                  'rhog2_cl':{'T0':{},'T1':{},'T2':{},'T3':{},'T4':{}},
                  'P':{'T1':{},'T2':{},'T3':{},'T4':{}},
                  'P_cl':{'T1':{},'T2':{},'T3':{},'T4':{}},
                  'mfrac':{'T1':{},'T2':{},'T3':{},'T4':{}},
                  'header':{}
                  }

    while 1 :
        clusterheader = input.readline()
        if ( clusterheader ) :
            header = load_clusterheader( clusterheader )
            id = header['id']

            if ( clusters == None or cluster_hash.has_key(id) ) :
                line = input.readline()
                cols = line.split()
                # Read in first bin of profiles data for this cluster id
                profiles = zeros( (num_bins,len(cols)) )
                profiles[0,:] = array([ float(col) for col in cols ])
                print len(cols)
                # Read in subsequent bins of profiles data for this cluster id
                for i in xrange(num_bins-1) :
                    line = input.readline()
                    profiles[i+1,:] = array([ float(x) for x in line.split() ])

                ( hclumppro['binr'][id], hclumppro['rr'][id],
                  hclumppro['Tmw'][id], hclumppro['Tsp'][id],
                  hclumppro['rhog']['T0'][id], hclumppro['rhog2_cl']['T0'][id],
                  hclumppro['rhog']['T1'][id], hclumppro['rhog2_cl']['T1'][id],
                  hclumppro['rhog']['T2'][id], hclumppro['rhog2_cl']['T2'][id],
                  hclumppro['rhog']['T3'][id], hclumppro['rhog2_cl']['T3'][id],
                  hclumppro['rhog']['T4'][id], hclumppro['rhog2_cl']['T4'][id],
                  hclumppro['P']['T1'][id], hclumppro['P_cl']['T1'][id],
                  hclumppro['P']['T2'][id], hclumppro['P_cl']['T2'][id],
                  hclumppro['P']['T3'][id], hclumppro['P_cl']['T3'][id],
                  hclumppro['P']['T4'][id], hclumppro['P_cl']['T4'][id],
                  hclumppro['mfrac']['T1'][id], hclumppro['mfrac']['T2'][id],
                  hclumppro['mfrac']['T3'][id], hclumppro['mfrac']['T4'][id]
                  ) = profiles[:,:].transpose()
                               
                # Populate header outputs into hclumppro
                hclumppro['header'][id] = header
                del profiles, header
                
            else :
                for i in xrange(num_bins) :
                    input.readline()
        else :
            break

    input.close()

    return hclumppro


def load_bzpro_profile(bzpro,clusters=None,num_bins=100) :
    '''
    Reads in hbzpro (metallicity profiles)
    Returns dictionary of dictionaries:
        hzpro = { 'binr':{}, 'rr':{}, [/h kpc]
              'Zg_II':{}, 'Zg_Ia':{}, 'Zhg_II':{}, 'Zhg_Ia':{}, (hot gas:T>1e6K)
              'Zsl_II':{}, 'Zsl_Ia':{}, 'Zst_II':{}, 'Zst_Ia':{},
              'Zstnew_II':{}, 'Zstnew_Ia':{}, 'Zg_IImass':{}, 'Zg_Iamass':{},
              'Zhg_IImass':{}, 'Zhg_Iamass':{}, 'Zsl_IImass':{}, 'Zsl_Iamass':{},
              'tst':{}, [Gyr]
              'header':{}
              }
    Zg are the metallicities of the gas, Zsl is the spectroscopically weighted metallicities in the gas, Z* is the metallicity contained in the stars that have not yet been released
    
    Ex: 
    hbzpro=load_bzpro_profile(h_bzpro_a?.????.dat,clusters=[1,2],num_bins=99)
    Zg_Ia=hzpro['Zg_Ia'][1] # Returns Z_I profile (array) of halo 1
    '''    
    input = open(bzpro,'r')

    if ( clusters != None ) :
        cluster_hash = {}
        for c in clusters :
            cluster_hash[c] = 1

    while 1 : # Skip over header lines
        line = input.readline()
        if line[0:2] == '##' :
            break

    hzpro = { 'binr':{}, 'rr':{},
              'Zg_II':{}, 'Zg_Ia':{}, 'Zhg_II':{}, 'Zhg_Ia':{},
              'Zsl_II':{}, 'Zsl_Ia':{}, 'Zst_II':{}, 'Zst_Ia':{},
              'Zstnew_II':{}, 'Zstnew_Ia':{}, 
              'tst':{}, 'header':{}
              }
    while 1 :
        clusterheader = input.readline()

        if ( clusterheader ) :
            header = load_clusterheader( clusterheader )
            id = header['id']

            if ( clusters == None or cluster_hash.has_key(id) ) :
                line = input.readline()
                cols = line.split()
                # Read in first bin of profiles data for this cluster id
                profiles = zeros( (num_bins,len(cols)) )
                profiles[0,:] = array([ float(col) for col in cols ])

                # Read in subsequent bins of profiles data for this cluster id
                for i in xrange(num_bins-1) :
                    line = input.readline()
                    profiles[i+1,:] = array([ float(x) for x in line.split() ])
                    
                    ( hzpro['binr'][id], hzpro['rr'][id],
                      hzpro['Zg_II'][id], hzpro['Zg_Ia'][id],
                      hzpro['Zhg_II'][id], hzpro['Zhg_Ia'][id], 
                      hzpro['Zsl_II'][id], hzpro['Zsl_Ia'][id], 
                      hzpro['Zst_II'][id], hzpro['Zst_Ia'][id], 
                      hzpro['Zstnew_II'][id], hzpro['Zstnew_Ia'][id],
                      hzpro['tst'][id] ) = profiles[:,:].transpose()

                # Populate header outputs into hzpro
                hzpro['header'][id] = header
                del profiles, header

            else :
                for i in xrange(num_bins) :
                    input.readline()
        else :
            break
    input.close()

    return hzpro

def load_differential_bmpro_profile(bmpro,clusters=None,num_bins=100) :
    '''
    Reads in hbmpro (mass profiles from clump analysis)
    Returns dictionary of dictionaries: 
    hbmdiffpro = { 'binr':{}, 'rr':{}, 'Mdm':{}, 'Mg':{},
                'Mcg':{}, 'Mst':{}, 'Mst_init':{},
                'Mst_new':{}, 'dMdm':{}, 'dMg':{}, 'dMhg':{},
                'dMst':{}, 'dMst_init':{}, 'dMst_new':{},
                'dMmetalgas_II':{}, 'dMmetalgas_Ia':{},
                'dMmetalhg_II':{}, 'dMmetalhg_Ia':{},
                'dMmetalst_II':{}, 'dMmetalst_Ia':{},
                'header':{}
                }
    Ex: 
    hbmdiffpro=load_differential_bmpro_profile(h_bmpro_a?.????.dat,clusters=[1,2],num_bins=99)
    dMcg=hbmdiffpro['dMhg'][1] # Returns differential hot gas mass profile (array) of halo 1 - differential mass profile is useful in calculating Yeff, etc.    
    '''
    input = open(bmpro,'r')
    if ( clusters != None ) :
        cluster_hash = {}
        for c in clusters :
            cluster_hash[c] = 1

    while 1 : # Skip over header lines
        line = input.readline()
        if line[0:2] == '##' :
            break

    hbmdiffpro = { 'binr':{}, 'Mdm':{}, 'Mg':{},
                   'Mcg':{}, 'Mst':{}, 'Mst_init':{},
                   'Mst_new':{}, 'dMdm':{}, 'dMg':{}, 'dMhg':{},
                   'dMst':{}, 'dMst_init':{}, 'dMst_new':{},
                   'dMmetalgas_II':{}, 'dMmetalgas_Ia':{},
                   'dMmetalhg_II':{}, 'dMmetalhg_Ia':{},
                   'dMmetalst_II':{}, 'dMmetalst_Ia':{},
                   'header':{}
                   }
    while 1 :
        clusterheader = input.readline()
        if ( clusterheader ) :
            header = load_clusterheader( clusterheader )
            id = header['id']

            if ( clusters == None or cluster_hash.has_key(id) ) :
                line = input.readline()
                cols = line.split()

                # Read in first bin of profiles data for this cluster id
                profiles = zeros( (num_bins,len(cols)) )
                profiles[0,:] = array([ float(col) for col in cols ])

                # Read in subsequent bins of profiles data for this cluster id
                for i in xrange(num_bins-1) :
                    line = input.readline()
                    profiles[i+1,:] = array([ float(x) for x in line.split() ])
                    
                    ( hbmdiffpro['binr'][id], hbmdiffpro['rr'][id],
                      hbmdiffpro['Mdm'][id], hbmdiffpro['Mg'][id],
                      hbmdiffpro['Mcg'][id], hbmdiffpro['Mst'][id], 
                      hbmdiffpro['Mst_init'][id], hbmdiffpro['Mst_new'][id], 
                      hbmdiffpro['dMdm'][id], hbmdiffpro['dMg'][id], 
                      hbmdiffpro['dMhg'][id], hbmdiffpro['dMst'][id],
                      hbmdiffpro['dMst_init'][id], hbmdiffpro['dMst_new'][id], 
                      hbmdiffpro['dMmetalgas_II'][id], hbmdiffpro['dMmetalgas_Ia'][id],
                      hbmdiffpro['dMmetalhg_II'][id], hbmdiffpro['dMmetalhg_Ia'][id],
                      hbmdiffpro['dMmetalst_II'][id], hbmdiffpro['dMmetalst_Ia'][id]
                      ) = profiles[:,:].transpose()

                # Populate header outputs into hbmdiffpro
                hbmdiffpro['header'][id] = header
                del profiles, header
                
            else :
                for i in xrange(num_bins) :
                    input.readline()
        else :
            break
    input.close()

    return hbmdiffpro


                
                
        
    
                        
                   
    

